import tkinter as tk
from tkinter import ttk, filedialog
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from matplotlib.ticker import FuncFormatter
import numpy as np
from spleeter.separator import Separator
import parselmouth
import os

# 初始化分离器
separator = Separator('spleeter:2stems')

# 设置 x 轴刻度的显示格式
def format_time(value, _):
    value = int(value)
    m = value // (100 * 60)
    s = (value % (100 * 60)) / 100
    return f"{m}:{s:.1f}"

class ScatterPlotApp:
    def __init__(self, root):
        self.root = root
        self.root.title('人声对比器')

        # Set up GUI components
        self.create_widgets()

    def create_widgets(self):
        # Frame for controls
        control_frame = ttk.Frame(self.root)
        control_frame.pack(pady=10)

        # Data 1: File selection button
        data1_frame = ttk.Frame(control_frame)
        data1_frame.pack(side='top', pady=5)
        ttk.Label(data1_frame, text='歌曲1:').pack(side='left', padx=5)
        ttk.Button(data1_frame, text="选择文件", command=self.load_data1).pack(side='left', padx=10)

        # Data 2: File selection button
        data2_frame = ttk.Frame(control_frame)
        data2_frame.pack(side='top', pady=5)
        ttk.Label(data2_frame, text='歌曲2:').pack(side='left', padx=5)
        ttk.Button(data2_frame, text="选择文件", command=self.load_data2).pack(side='left', padx=10)

        # X Offset slider
        self.x_offset_var = tk.DoubleVar(value=0)
        x_offset_frame = ttk.Frame(control_frame)
        x_offset_frame.pack(side='top', pady=5)
        ttk.Label(x_offset_frame, text='位移:').pack(side='left', padx=5)
        x_offset_slider = ttk.Scale(x_offset_frame, from_=0, to=1, variable=self.x_offset_var, orient='horizontal', length=200, name="x_offset_slider")
        x_offset_slider.pack(side='left', padx=10)
        x_offset_slider.bind("<B1-Motion>", self.update_plot)
        x_offset_slider.bind("<ButtonRelease-1>", self.update_plot_click)

        # X Scale slider
        self.x_scale_var = tk.DoubleVar(value=0.01)
        x_scale_frame = ttk.Frame(control_frame)
        x_scale_frame.pack(side='top', pady=5)
        ttk.Label(x_scale_frame, text='缩放:').pack(side='left', padx=5)
        x_scale_slider = ttk.Scale(x_scale_frame, from_=0.001, to=0.1, variable=self.x_scale_var, orient='horizontal', length=200, name="x_scale_slider")
        x_scale_slider.pack(side='left', padx=10)
        x_scale_slider.bind("<B1-Motion>", self.update_plot)
        x_scale_slider.bind("<ButtonRelease-1>", self.update_plot_click)

        # X Offset 2 slider
        self.x_offset_2_var = tk.DoubleVar(value=0)
        x_offset_2_frame = ttk.Frame(control_frame)
        x_offset_2_frame.pack(side='top', pady=5)
        ttk.Label(x_offset_2_frame, text='偏移:').pack(side='left', padx=5)
        x_offset_2_spinbox = ttk.Spinbox(x_offset_2_frame, from_=-10000, to=10000, textvariable=self.x_offset_2_var)
        x_offset_2_spinbox.pack(side='left', padx=10)
        self.x_offset_2_var.trace_add("write", self.on_spinbox_change)

        # Scatter plot canvas
        self.fig, self.ax = plt.subplots(figsize=(6, 4))
        self.scatter_plot = FigureCanvasTkAgg(self.fig, master=self.root)
        self.scatter_plot.get_tk_widget().pack()

    def load_data1(self):
        file_path = filedialog.askopenfilename(title="选择歌曲1")
        print(file_path)
        if file_path:
            folder = file_path.split('/')[-1].split('.')[0]
            vocal_file = f'output_path/{folder}/vocals.wav'
            if not os.path.exists(vocal_file):
                # 进行音频源分离
                print('正在分离人声...')
                separator.separate_to_file(file_path, 'output_path')

            # 读取音频文件
            snd = parselmouth.Sound(vocal_file)

            # 提取基频信息
            print('正在提取基频...')
            pitch = snd.to_pitch_ac()
            pitch_values = pitch.selected_array['frequency']

            # Save data to instance variable
            self.data1 = pitch_values

            # Update scatter plot
            self.update_plot(None)

    def load_data2(self):
        file_path = filedialog.askopenfilename(title="选择歌曲2")
        print(file_path)
        if file_path:
            folder = file_path.split('/')[-1].split('.')[0]
            vocal_file = f'output_path/{folder}/vocals.wav'
            if not os.path.exists(vocal_file):
                # 进行音频源分离
                print('正在分离人声...')
                separator.separate_to_file(file_path, 'output_path')

            # 读取音频文件
            snd = parselmouth.Sound(vocal_file)

            # 提取基频信息
            print('正在提取基频...')
            pitch = snd.to_pitch_ac()
            pitch_values = pitch.selected_array['frequency']

            # Save data to instance variable
            self.data2 = pitch_values

            # Update scatter plot
            self.update_plot(None)

    def on_spinbox_change(self, *args):
        self.update_plot(None)

    def update_plot(self, event):
        if hasattr(self, 'data1'):
            x_offset = self.x_offset_var.get()
            x_scale = self.x_scale_var.get()
            try:
                x_offset_2 = self.x_offset_2_var.get()
            except:
                x_offset_2 = 0

            # Update x-axis limits for Data 1
            l1 = len(self.data1)
            x_min1 = int(l1 * (1 - x_scale) * x_offset)
            x_max1 = min(x_min1 + int(l1 * x_scale), l1)

            # Update x-axis limits for Data 2
            x_min2 = int(x_min1 - x_offset_2)
            x_max2 = int(x_max1 - x_offset_2)

            # Update scatter plot for Data 1
            self.ax.clear()
            non_zero_indices_1 = (self.data1 != 0) & (
                    x_min1 <= np.arange(len(self.data1))) & (
                    np.arange(len(self.data1)) <= x_max1)
            non_zero_data_1 = self.data1[non_zero_indices_1]
            x1 = np.where(non_zero_indices_1)[0]
            self.ax.scatter(x1, non_zero_data_1, color='blue', marker='o', label='Data 1', s=1)

            if hasattr(self, 'data2'):
                non_zero_indices_2 = (self.data2 != 0) & (x_min2 <= np.arange(len(self.data2))) & (np.arange(len(self.data2)) <= x_max2)
                non_zero_data_2 = self.data2[non_zero_indices_2]
                x2 = np.where(non_zero_indices_2)[0]
                x2 = x2 + x_offset_2
                self.ax.scatter(x2, non_zero_data_2, color='red', marker='o', label='Data 2', s=1)

            self.ax.set_xlabel('Index')
            self.ax.set_ylabel('Y-axis Label')
            self.ax.set_xlim([x_min1, x_max1])
            self.ax.xaxis.set_major_formatter(FuncFormatter(format_time))

            self.scatter_plot.draw()


    def update_plot_click(self, event):
        if event.widget.winfo_name() == "x_offset_slider":
            click_position = event.x / 200.0  # Assuming the length of slider is 200 pixels
            self.x_offset_var.set(click_position)
        elif event.widget.winfo_name() == "x_scale_slider":
            click_position = event.x / 200.0 * (0.1 - 0.001) + 0.001
            self.x_scale_var.set(click_position)
        self.update_plot(None)


if __name__ == "__main__":
    root = tk.Tk()
    app = ScatterPlotApp(root)
    root.mainloop()
